package main

import (
	"context"
	"golang.org/x/crypto/ssh"
	"io"
	"log"
	"net"
	"ssh-tunnel/util"
	"sync"
	"sync/atomic"
	"time"
)

type SSHTunnel struct {
	Server            string
	User              string
	Password          string
	Timeout           int
	RetryInterval     int
	KeepAliveInterval int
	KeepAliveCountMax int

	ctx context.Context
	cli *ssh.Client
}

func NewSSHTunnel() *SSHTunnel {
	r := &SSHTunnel{
		Server:            *server,
		User:              *user,
		Password:          *password,
		Timeout:           *timeout,
		RetryInterval:     *retry,
		KeepAliveInterval: 30,
		KeepAliveCountMax: 5,
	}
	ctx, cancelFunc := context.WithCancel(context.Background())
	r.ctx = ctx

	registerOnShutdown(func() {
		cancelFunc()
	})

	go r.sshClient()
	return r
}

func (r *SSHTunnel) sshClient() {
	for {
		select {
		case <-r.ctx.Done():
			return
		default:
			r.sshClient0()
			<-time.After(time.Duration(r.RetryInterval) * time.Second)
			log.Printf("ssh connection retrying...\n")
		}
	}
}

func (r *SSHTunnel) sshClient0() {
	var once sync.Once
	config := &ssh.ClientConfig{
		Timeout: time.Duration(r.Timeout) * time.Second,
		User:    r.User,
		Auth: []ssh.AuthMethod{
			ssh.Password(r.Password),
			ssh.PasswordCallback(func() (string, error) {
				return "", nil
			}),
			ssh.KeyboardInteractive(func(user, instruction string, questions []string, echos []bool) (answers []string, err error) {
				log.Println(`user: ` + user)
				log.Println(`Instruction: ` + instruction)
				log.Println(`Questions:`)
				for q := range questions {
					log.Printf(`%d\n`, q)
				}
				countQuestions := len(questions)
				if countQuestions == 1 {
					answers = make([]string, countQuestions, countQuestions)
					answers[0] = ""
				} else if countQuestions > 1 {
					log.Println(`The ssh server is asking multiple questions! s program cannot handle s case.`)
				}
				err = nil
				return
			}),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Config: ssh.Config{
			Ciphers: []string{
				"aes128-ctr",
				"aes192-ctr",
				"aes256-ctr",
				"aes128-gcm@openssh.com",
				"arcfour256",
				"arcfour128",
				"aes128-cbc",
				"3des-cbc",
				"aes192-cbc",
				"aes256-cbc",
			},
			KeyExchanges: []string{
				"curve25519-sha256@libssh.org",
				"ecdh-sha2-nistp256",
				"ecdh-sha2-nistp384",
				"ecdh-sha2-nistp521",
				"diffie-hellman-group14-sha256",
				"diffie-hellman-group14-sha1",
				"diffie-hellman-group1-sha1",
			},
			MACs: []string{
				"hmac-sha2-256",
				"hmac-sha2-512",
				"hmac-sha1",
				"hmac-sha1-96",
				"hmac-md5",
				"hmac-md5-96",
			},
		},
	}
	log.Printf("ssh connecting %s...\n", r.Server)
	// Connect to the server host via SSH.
	cli, err := ssh.Dial("tcp", r.Server, config)
	if err != nil {
		once.Do(func() { log.Printf("ssh dial error: %v", err) })
		return
	}
	defer cli.Close()
	r.cli = cli
	log.Printf("ssh connected %s\n", r.Server)
	registerOnShutdown(func() {
		log.Printf("ssh disconnect %s...\n", r.Server)
		cli.Close()
	})
	r.keepAliveMonitor(&once)
}

func (r *SSHTunnel) StartTunnel(localAddr string, remoteAddr string) {
	l, err := net.Listen("tcp", localAddr)
	if err != nil {
		log.Println(err.Error())
		return
	}
	defer l.Close()
	registerOnShutdown(func() {
		l.Close()
	})
	log.Printf("L:%s <==> R:%s\n", localAddr, remoteAddr)
	for {
		lConn, err := l.Accept()
		if err != nil {
			continue
		}
		if r.cli == nil {
			continue
		}
		go func() {
			rConn, err := r.cli.Dial("tcp", remoteAddr)
			if err != nil {
				log.Println(err.Error())
				return
			}
			defer rConn.Close()
			util.IOTransport(lConn, rConn)
		}()
	}
}

func (r *SSHTunnel) keepAliveMonitor(once *sync.Once) {
	if r.KeepAliveInterval == 0 || r.KeepAliveCountMax == 0 {
		return
	}
	// Detect when the ssh connection is closed.
	wait := make(chan error, 1)
	go func() {
		wait <- r.cli.Wait()
	}()

	// Repeatedly check if the remote server is still alive.
	var aliveCount int32
	ticker := time.NewTicker(time.Duration(r.KeepAliveInterval) * time.Second)
	defer ticker.Stop()
	for {
		select {
		case err := <-wait:
			if err != nil && err != io.EOF {
				once.Do(func() { log.Printf("ssh error: %v", err) })
			}
			return
		case <-ticker.C:
			if n := atomic.AddInt32(&aliveCount, 1); n > int32(r.KeepAliveCountMax) {
				once.Do(func() { log.Printf("ssh keep-alive termination") })
				r.cli.Close()
				return
			}
		}

		go func() {
			_, _, err := r.cli.SendRequest("keepalive@openssh.com", true, nil)
			if err == nil {
				atomic.StoreInt32(&aliveCount, 0)
			}
		}()
	}
}
